package org.fujay.langchain4jserver.infrastructure.persistence.repository;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import lombok.RequiredArgsConstructor;
import org.fujay.langchain4jserver.domain.knowledge.aggregate.KnowledgeBaseAggregate;
import org.fujay.langchain4jserver.domain.knowledge.entity.File;
import org.fujay.langchain4jserver.domain.knowledge.entity.KnowledgeBase;
import org.fujay.langchain4jserver.domain.knowledge.repository.FileRepository;
import org.fujay.langchain4jserver.domain.knowledge.repository.KnowledgeBaseRepository;
import org.fujay.langchain4jserver.domain.knowledge.vo.FileId;
import org.fujay.langchain4jserver.domain.knowledge.vo.KnowledgeBaseId;
import org.fujay.langchain4jserver.domain.knowledge.vo.KnowledgeBaseStatus;
import org.fujay.langchain4jserver.domain.model.vo.ModelId;
import org.fujay.langchain4jserver.infrastructure.persistence.assembler.KnowledgeAssembler;
import org.fujay.langchain4jserver.infrastructure.persistence.mybatis.entity.KnowledgeBasePo;
import org.fujay.langchain4jserver.infrastructure.persistence.mybatis.mapper.KnowledgeBaseMapper;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 知识库仓储实现
 */
@Repository
@RequiredArgsConstructor
public class KnowledgeBaseRepositoryImpl implements KnowledgeBaseRepository {

    private final KnowledgeBaseMapper knowledgeBaseMapper;
    private final FileRepository fileRepository;

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void save(KnowledgeBaseAggregate aggregate) {
        KnowledgeBasePo kbPO = KnowledgeAssembler.convertKbToPO(aggregate.getKnowledgeBase());
        if (knowledgeBaseMapper.selectById(kbPO.getId()) == null) {
            knowledgeBaseMapper.insert(kbPO);
        } else {
            knowledgeBaseMapper.updateById(kbPO);
        }
        List<File> filesToSave = aggregate.getFiles().values().stream().toList();
        if (!CollectionUtils.isEmpty(filesToSave)) {
            fileRepository.saveAll(filesToSave, aggregate.getId());
        }
    }

    @Override
    public Optional<KnowledgeBaseAggregate> findById(KnowledgeBaseId id) {
        KnowledgeBasePo kbPO = knowledgeBaseMapper.selectById(id.value());
        if (kbPO == null) {
            return Optional.empty();
        }
        // 加载关联的文件
        List<File> files = fileRepository.findByKnowledgeBaseId(id);
        Map<FileId, File> fileMap =
                files.stream().collect(Collectors.toMap(File::getId, Function.identity()));

        KnowledgeBase kb = KnowledgeAssembler.convertKbToDomain(kbPO);
        return Optional.of(KnowledgeBaseAggregate.load(kb, fileMap));
    }

    @Override
    public Optional<KnowledgeBaseAggregate> findByIdShallow(KnowledgeBaseId id) {
        KnowledgeBasePo kbPO = knowledgeBaseMapper.selectById(id.value());
        if (kbPO == null) {
            return Optional.empty();
        }
        KnowledgeBase kb = KnowledgeAssembler.convertKbToDomain(kbPO);
        return Optional.of(KnowledgeBaseAggregate.load(kb, Map.of()));
    }

    @Override
    public Optional<ModelId> findEmbeddingModelId(KnowledgeBaseId id) {

        KnowledgeBasePo kbPO = knowledgeBaseMapper.selectById(id.value());
        if (kbPO == null) {
            return Optional.empty();
        }
        return Optional.of(new ModelId(kbPO.getEmbeddingModelId()));
    }

    @Override
    public Optional<KnowledgeBaseAggregate> findByIdAndFile(KnowledgeBaseId id, FileId fileId) {

        KnowledgeBasePo kbPO = knowledgeBaseMapper.selectById(id.value());
        File file = fileRepository.findById(fileId).orElse(null);

        return Optional.empty();
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void deleteById(KnowledgeBaseId id) {
        // 1. 删除关联文件 (委托给 FileRepository)
        fileRepository.deleteByKnowledgeBaseId(id);
        // 2. 删除知识库自身
        knowledgeBaseMapper.deleteById(id.value());
    }

    @Override
    public List<KnowledgeBaseAggregate> findAllShallow() {
        // 只查询 KnowledgeBasePO，不加载文件
        return knowledgeBaseMapper.selectList(Wrappers.emptyWrapper()).stream()
                .map(KnowledgeAssembler::convertKbToDomain)
                .map(kb -> KnowledgeBaseAggregate.load(kb, Map.of())) // 创建不包含文件的聚合
                .collect(Collectors.toList());
    }

    @Override
    public List<KnowledgeBaseAggregate> findCompleted() {
        LambdaQueryWrapper<KnowledgeBasePo> queryWrapper =
                Wrappers.<KnowledgeBasePo>lambdaQuery()
                        .eq(KnowledgeBasePo::getStatus, KnowledgeBaseStatus.INDEXING_COMPLETED.getCode());
        return knowledgeBaseMapper.selectList(queryWrapper).stream()
                .map(KnowledgeAssembler::convertKbToDomain)
                .map(kb -> KnowledgeBaseAggregate.load(kb, Map.of()))
                .collect(Collectors.toList());
    }

    @Override
    public boolean existsByIndexName(String indexName) {
        return knowledgeBaseMapper.exists(
                Wrappers.<KnowledgeBasePo>lambdaQuery().eq(KnowledgeBasePo::getIndexName, indexName));
    }
}
